import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import itertools
import math
from sklearn.metrics import confusion_matrix
import module


def plot_confusion_matrix(cm,
                          target_names,
                          title='Confusion matrix',
                          cmap=plt.cm.Greens,#这个地方设置混淆矩阵的颜色主题，这个主题看着就干净~
                          normalize=True):
   
 
    accuracy = np.trace(cm) / float(np.sum(cm))
    misclass = 1 - accuracy

    if cmap is None:
        cmap = plt.get_cmap('Blues')

    plt.figure()
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()

    if target_names is not None:
        tick_marks = np.arange(len(target_names))
        plt.xticks(tick_marks, target_names, rotation=45)
        plt.yticks(tick_marks, target_names)

    if normalize:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]


    thresh = cm.max() / 1.5 if normalize else cm.max() / 2
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        if normalize:
            plt.text(j, i, "{:0.4f}".format(cm[i, j]),
                     horizontalalignment="center",
                     color="white" if cm[i, j] > thresh else "black")
        else:
            plt.text(j, i, "{:,}".format(cm[i, j]),
                     horizontalalignment="center",
                     color="white" if cm[i, j] > thresh else "black")


    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label\naccuracy={:0.4f}; misclass={:0.4f}'.format(accuracy, misclass))
    #这里这个savefig是保存图片，如果想把图存在什么地方就改一下下面的路径，然后dpi设一下分辨率即可。
	#plt.savefig('/content/drive/My Drive/Colab Notebooks/confusionmatrix32.png',dpi=350)
# 显示混淆矩阵
def plot_confuse(model, x_val, y_val):
    predictions = model.predict(x_val).argmax(axis=-1)
    conf_mat = confusion_matrix(y_true=y_val, y_pred=predictions)
    print(conf_mat)#输出混淆矩阵
    plot_confusion_matrix(conf_mat, normalize=False, target_names= label_names, title='Confusion Matrix')

def plot_error(model, x_val, y_val):
    predictions = model.predict(x_val).argmax(axis=-1).astype('int32')
    err_index = [i for i,x in enumerate(predictions==y_val) if not x]
    err_index = err_index[:10]
    plt.figure()
    for i,index in enumerate(err_index):
        plt.subplot(4,5, math.floor(i/5)*10+i%5+1)
        plt.axis('off')
        plt.imshow(x_val[index])
        plt.subplot(4,5,math.floor(i/5)*10+i%5+5+1)
        plt.axis('off')
        plt.text(0.,.5,'pred: {} \ntrue: {}'.format(label_names[predictions[index]],label_names[y_val[index]]))
# =========================================================================================
# 最后调用这个函数即可。 test_x是测试数据，test_y是测试标签（这里用的是One——hot向量）
# labels是一个列表，存储了类别的名字，最后会显示在横纵轴上。

_,_,x_test,y_test,label_names = module.load_data()
# # 测试集
# x_test,y_test = load_CIFAR_batch('test')
# #label_names
# label_names = load_pickle('meta')['fine_label_names']
model = tf.keras.models.load_model(module.type+'_cs_model.h5',custom_objects={'precision': module.precision,'recall':module.recall})
# model.compile(optimizer=tf.keras.optimizers.SGD(lr=0.1),loss=tf.keras.losses.categorical_crossentropy, metrics=['accuracy',recall,precision])
model.summary()
model.evaluate(x_test,y_test)
plot_confuse(model, x_test, y_test.argmax(axis=-1))
plot_error(model, x_test, y_test.argmax(axis=-1))
plt.show()# 显示混淆矩阵

